import easydict as edict
import json
import pickle
import random
import os
import os.path
from text2vec import SBert
import numpy as np

os.environ["CUDA_VISIBLE_DEVICES"] = '1'


def compute_emb(sentences):
    model = SBert('paraphrase-multilingual-MiniLM-L12-v2')
    embeddings = model.encode(sentences)
    return embeddings


def read_json(file_path):
    f = open(file_path, encoding="utf-8")
    file = json.load(f)
    print(type(file))
    image_list = file['images']
    annotation_list = file['annotations']
    return image_list, annotation_list

split = 'val'
images, annotations = read_json('data/captions_train-val2014/annotations/captions_val2014.json')
data_dir = 'data'
split_dir = os.path.join(data_dir, split)
image_dir = os.path.join(data_dir, 'val2014')
image_name = os.listdir(image_dir)
print(image_name)
filenames = []
captions = []
print('start')

num = 0
for img in images:
    print('hit = ', num)
    if img['file_name'] in image_name:
        num += 1
        filenames.append(img['file_name'])
        # img_embedding = []
        img_caption = []
        img_id = img['id']
        for anno in annotations:
            if anno['image_id'] == img_id:
                print('hit in ', img_id)
                img_caption.append(anno['caption'])
        img_caption = img_caption[random.randint(0, len(img_caption) - 1)]
        captions.append(img_caption)
        # embeddings.append(compute_emb(img_caption))
        # print(len(embeddings))

embeddings = compute_emb(captions)
print(type(captions), np.shape(captions))
print(captions)
print(type(embeddings), np.shape(embeddings))
print(embeddings)

filenames_path = os.path.join(split_dir, 'filenames.pkl')
embeddings_path = os.path.join(split_dir, 'embeddings.pkl')
with open(filenames_path, 'wb') as ff:
    pickle.dump(filenames, ff)
with open(embeddings_path, 'wb') as fe:
    pickle.dump(embeddings, fe)
